// main.go
package main

import (
	"alan/exp3/src/entities"
	"alan/exp3/src/system"
	"crypto/rand"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"math/big"
	"os"
	"time"
)

// PerformanceResult holds the timing data for one experiment run.
type PerformanceResult struct {
	Experiment     string
	Parameter      string
	Value          int
	EncryptTimeMs  float64
	AuditTimeMs    float64
	ProofGenTimeMs float64
	VerfyTimeMs    float64
}

// 基础实验
func basicExp() {
	logFile, err := os.OpenFile("run.log", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open log file: %v", err)
	}
	defer logFile.Close()
	mw := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(mw)

	log.Println("--- Step 1: System Setup ---")
	const K, FileSize, BlockSize, ChallengedBlocks = 5, 1024, 16, 10

	cspIDs := make([]string, K)
	for i := 0; i < K; i++ {
		cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
	}
	params, err := system.NewSystemParams(cspIDs)
	if err != nil {
		log.Fatalf("Failed to create system params: %v", err)
	}

	log.Println("--- Step 2: Key Generation and File Outsourcing ---")
	dataOwner, err := entities.NewDataOwner(params)
	if err != nil {
		log.Fatalf("Failed to create data owner: %v", err)
	}

	dummyFile := make([]byte, FileSize)
	rand.Read(dummyFile)
	encryptedDataMap, err := dataOwner.EncryptFile(dummyFile, "doc1", BlockSize)
	if err != nil {
		log.Fatalf("Failed to encrypt file: %v", err)
	}

	csps := make([]*entities.CSP, K)
	for i := 0; i < K; i++ {
		csps[i] = entities.NewCSP(cspIDs[i], params)
		csps[i].StoreFile("doc1", encryptedDataMap[cspIDs[i]])
	}
	log.Printf("All %d CSPs have stored their encrypted copies.", K)

	log.Println("\n--- Step 3: Normal Auditing Process ---")
	auditRequest, err := dataOwner.RequestAudit(0, ChallengedBlocks, "doc1")
	if err != nil {
		log.Fatalf("Failed to request audit: %v", err)
	}

	partialProofs := make([]*big.Int, K)
	for i, csp := range csps {
		proof, err := csp.ProofGen(auditRequest)
		if err != nil {
			log.Fatalf("CSP %s failed to generate proof: %v", csp.ID, err)
		}
		partialProofs[i] = proof
	}

	allVerifiedNormal := true
	for _, csp := range csps {
		completeProof := csp.AggregateProof(partialProofs)
		if !csp.VerifyProof(auditRequest, completeProof) {
			allVerifiedNormal = false
		}
	}

	if allVerifiedNormal {
		log.Println("SUCCESS: All CSPs passed the verification in the normal process.")
	} else {
		log.Println("FAILURE: Verification failed in the normal process.")
	}

	log.Println("\n--- Step 4: Auditing Process After Data Tampering ---")
	// 重新生成一个审计请求，以增加命中被篡改块的概率
	log.Println("Data Owner is generating a NEW, RANDOM audit request for the tampering check.")
	tamperedAuditRequest, err := dataOwner.RequestAudit(0, ChallengedBlocks, "doc1")
	if err != nil {
		log.Fatalf("Failed to request tampered audit: %v", err)
	}

	csps[1].EncryptedFile[5].Add(csps[1].EncryptedFile[5], big.NewInt(1))
	log.Printf("!!! Tampering data on %s, block %d !!!", csps[1].ID, 5)

	tamperedPartialProofs := make([]*big.Int, K)
	for i, csp := range csps {
		proof, _ := csp.ProofGen(tamperedAuditRequest)
		tamperedPartialProofs[i] = proof
	}

	detectedTampering := false
	for _, csp := range csps {
		completeProof := csp.AggregateProof(tamperedPartialProofs)
		if !csp.VerifyProof(tamperedAuditRequest, completeProof) {
			detectedTampering = true
		}
	}
	if detectedTampering {
		log.Println("SUCCESS: Data tampering was successfully detected.")
	} else {
		log.Println("FAILURE: No CSP detected the data tampering. (This can happen by chance if the tampered block was not challenged. Run again.)")
	}
}

// 性能实验
// 实验1：文件大小
func fileSizeExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Running Experiment: File Size Variation >>>>>>>>>>")
	const K, BlockSize = 5, 2 * 1024
	fileSizes := []int{2, 4, 6, 8, 10, 12, 14, 16, 18} // 文件大小: 2MB to 18MB

	for _, sizeMB := range fileSizes {
		fileSize := sizeMB * 1024 * 1024
		log.Printf("\n--- Testing File Size: %d MB ---", sizeMB)

		cspIDs := make([]string, K)
		for i := range K {
			cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
		}
		params, _ := system.NewSystemParams(cspIDs)
		dataOwner, _ := entities.NewDataOwner(params)

		dummyFile := make([]byte, fileSize)
		rand.Read(dummyFile)

		// 加密计时
		startEncrypt := time.Now()
		encryptedDataMap, _ := dataOwner.EncryptFile(dummyFile, "perf_doc", BlockSize)
		encryptTime := time.Since(startEncrypt)

		csps := make([]*entities.CSP, K)
		for i := range K {
			csps[i] = entities.NewCSP(cspIDs[i], params)
			csps[i].StoreFile("perf_doc", encryptedDataMap[cspIDs[i]])
		}

		numBlocks := fileSize / BlockSize
		if fileSize%BlockSize != 0 {
			numBlocks++
		}
		ChallengedBlocks := numBlocks

		// 审计请求计时
		startAudit := time.Now()
		auditRequest, _ := dataOwner.RequestAudit(0, ChallengedBlocks, "perf_doc")
		auditTime := time.Since(startAudit)

		// 所有部分证明生成总计时->平均时间
		var totalProofGenTime time.Duration
		partialProofs := make([]*big.Int, K)
		for k, csp := range csps {
			startProofGen := time.Now()
			proof, _ := csp.ProofGen(auditRequest)
			totalProofGenTime += time.Since(startProofGen)
			partialProofs[k] = proof
		}
		avgProofGenTime := totalProofGenTime / time.Duration(K)

		// 验证计时
		var totalVerifyTime time.Duration
		for _, csp := range csps {
			completeProof := csp.AggregateProof(partialProofs)
			startVerify := time.Now()
			csp.VerifyProof(auditRequest, completeProof)
			totalVerifyTime += time.Since(startVerify)
		}
		avgVerifyTime := totalVerifyTime / time.Duration(K)

		result := PerformanceResult{
			Experiment: "FileSize", Parameter: "FileSizeMB", Value: sizeMB,
			EncryptTimeMs:  float64(encryptTime.Microseconds()) / 1000.0,
			AuditTimeMs:    float64(auditTime.Microseconds()) / 1000.0,
			ProofGenTimeMs: float64(avgProofGenTime.Microseconds()) / 1000.0,
			VerfyTimeMs:    float64(avgVerifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.EncryptTimeMs), fmt.Sprintf("%.3f", result.AuditTimeMs),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs), fmt.Sprintf("%.3f", result.VerfyTimeMs),
		})
		writer.Flush()
	}
}

// 实验2：CSP数量
func cspCountExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Running Experiment: CSP Count Variation >>>>>>>>>>")
	const FileSize, BlockSize = 2 * 1024 * 1024, 2 * 1024
	cspCounts := []int{5, 7, 9, 11, 13, 15, 17, 19, 21}

	for _, k := range cspCounts {
		log.Printf("\n--- Testing CSP Count: %d ---", k)

		cspIDs := make([]string, k)
		for i := 0; i < k; i++ {
			cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
		}
		params, _ := system.NewSystemParams(cspIDs)
		dataOwner, _ := entities.NewDataOwner(params)

		dummyFile := make([]byte, FileSize)
		rand.Read(dummyFile)

		startEncrypt := time.Now()
		encryptedDataMap, _ := dataOwner.EncryptFile(dummyFile, "perf_doc", BlockSize)
		encryptTime := time.Since(startEncrypt)

		csps := make([]*entities.CSP, k)
		for i := 0; i < k; i++ {
			csps[i] = entities.NewCSP(cspIDs[i], params)
			csps[i].StoreFile("perf_doc", encryptedDataMap[cspIDs[i]])
		}

		numBlocks := FileSize / BlockSize
		if FileSize%BlockSize != 0 {
			numBlocks++
		}
		ChallengedBlocks := numBlocks

		startAudit := time.Now()
		auditRequest, _ := dataOwner.RequestAudit(0, ChallengedBlocks, "perf_doc")
		auditTime := time.Since(startAudit)

		var totalProofGenTime time.Duration
		partialProofs := make([]*big.Int, k)
		for k, csp := range csps {
			startProofGen := time.Now()
			proof, _ := csp.ProofGen(auditRequest)
			totalProofGenTime += time.Since(startProofGen)
			partialProofs[k] = proof
		}
		avgProofGenTime := totalProofGenTime / time.Duration(k)

		var totalVerifyTime time.Duration
		for _, csp := range csps {
			completeProof := csp.AggregateProof(partialProofs)
			startVerify := time.Now()
			csp.VerifyProof(auditRequest, completeProof)
			totalVerifyTime += time.Since(startVerify)
		}
		avgVerifyTime := totalVerifyTime / time.Duration(k)

		result := PerformanceResult{
			Experiment: "CSPCount", Parameter: "NumCSPs", Value: k,
			EncryptTimeMs:  float64(encryptTime.Microseconds()) / 1000.0,
			AuditTimeMs:    float64(auditTime.Microseconds()) / 1000.0,
			ProofGenTimeMs: float64(avgProofGenTime.Microseconds()) / 1000.0,
			VerfyTimeMs:    float64(avgVerifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.EncryptTimeMs), fmt.Sprintf("%.3f", result.AuditTimeMs),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs), fmt.Sprintf("%.3f", result.VerfyTimeMs),
		})
		writer.Flush()
	}
}

// 实验3：分块大小
func blockSizeExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Running Experiment: Number of Blocks Variation >>>>>>>>>>")
	const K, FileSize = 5, 2 * 1024 * 1024
	// 我们仍然通过改变块大小来改变块数量
	blockSizes := []int{2, 4, 6, 8, 10, 12, 14, 16, 18} // in KB

	for _, sizeKB := range blockSizes {
		blockSize := sizeKB * 1024

		numBlocks := FileSize / blockSize
		if FileSize%blockSize != 0 {
			numBlocks++
		}

		log.Printf("\n--- Testing Block Size: %d KB (Number of Blocks: %d) ---", sizeKB, numBlocks)

		cspIDs := make([]string, K)
		for i := 0; i < K; i++ {
			cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
		}
		params, _ := system.NewSystemParams(cspIDs)
		dataOwner, _ := entities.NewDataOwner(params)

		dummyFile := make([]byte, FileSize)
		rand.Read(dummyFile)

		startEncrypt := time.Now()
		encryptedDataMap, _ := dataOwner.EncryptFile(dummyFile, "perf_doc", blockSize)
		encryptTime := time.Since(startEncrypt)

		csps := make([]*entities.CSP, K)
		for i := 0; i < K; i++ {
			csps[i] = entities.NewCSP(cspIDs[i], params)
			csps[i].StoreFile("perf_doc", encryptedDataMap[cspIDs[i]])
		}

		ChallengedBlocks := numBlocks

		startAudit := time.Now()
		auditRequest, _ := dataOwner.RequestAudit(0, ChallengedBlocks, "perf_doc")
		auditTime := time.Since(startAudit)

		var totalProofGenTime time.Duration
		partialProofs := make([]*big.Int, K)
		for k, csp := range csps {
			startProofGen := time.Now()
			proof, _ := csp.ProofGen(auditRequest)
			totalProofGenTime += time.Since(startProofGen)
			partialProofs[k] = proof
		}
		avgProofGenTime := totalProofGenTime / time.Duration(K)

		var totalVerifyTime time.Duration
		for _, csp := range csps {
			completeProof := csp.AggregateProof(partialProofs)
			startVerify := time.Now()
			csp.VerifyProof(auditRequest, completeProof)
			totalVerifyTime += time.Since(startVerify)
		}
		avgVerifyTime := totalVerifyTime / time.Duration(K)

		result := PerformanceResult{
			Experiment: "NumBlocks", Parameter: "NumberOfBlocks", Value: numBlocks, // 使用块数量作为参数和值
			EncryptTimeMs:  float64(encryptTime.Microseconds()) / 1000.0,
			AuditTimeMs:    float64(auditTime.Microseconds()) / 1000.0,
			ProofGenTimeMs: float64(avgProofGenTime.Microseconds()) / 1000.0,
			VerfyTimeMs:    float64(avgVerifyTime.Microseconds()) / 1000.0,
		}

		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.EncryptTimeMs), fmt.Sprintf("%.3f", result.AuditTimeMs),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs), fmt.Sprintf("%.3f", result.VerfyTimeMs),
		})
		writer.Flush()
	}
}

func main() {
	// Setup logging
	logFile, err := os.OpenFile("experiment.log", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open log file: %v", err)
	}
	defer logFile.Close()
	mw := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(mw)

	// Setup CSV writer
	csvFile, err := os.Create("performance_results.csv")
	if err != nil {
		log.Fatalf("Failed to create CSV file: %v", err)
	}
	defer csvFile.Close()
	csvWriter := csv.NewWriter(csvFile)
	defer csvWriter.Flush()
	csvWriter.Write([]string{"Experiment", "Parameter", "Value", "EncryptTimeMs", "AuditTimeMs", "ProofGenTimeMs", "VerifyTimeMs"})

	// run exp
	fileSizeExp(csvWriter)
	cspCountExp(csvWriter)
	blockSizeExp(csvWriter)

	log.Println("\n\n<<<<<<<<<< All experiments finished. Results saved to performance_results.csv >>>>>>>>>>")
}
